import torch
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,))])

train_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=True,
                               transform=transform,
                               download=True)
test_dataset = datasets.MNIST(root='../dataset/mnist',
                              train=False,
                              transform=transform,
                              download=True)
train_dataloader = DataLoader(dataset=train_dataset,
                              batch_size=64,
                              shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,
                             batch_size=64,
                             shuffle=False)


class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.channels = channels
        self.Conv_1 = nn.Conv2d(channels, channels, kernel_size=(3, 3), padding=1)
        self.Conv_2 = nn.Conv2d(channels, channels, kernel_size=(3, 3), padding=1)

    def forward(self, x):
        y = F.relu(self.Conv_1(x))
        y = self.Conv_2(y)
        y = F.relu(x + y)
        return y


class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.Conv_1 = nn.Conv2d(1, 16, kernel_size=(5, 5))
        self.Conv_2 = nn.Conv2d(16, 32, kernel_size=(5, 5))
        self.mp = nn.MaxPool2d(2)
        self.ResBlock_1 = ResidualBlock(16)
        self.ResBlock_2 = ResidualBlock(32)
        self.fc = nn.Linear(512, 10)

    def forward(self, x):
        batch_size = x.size(0)
        x = self.mp(F.relu(self.Conv_1(x)))
        x = self.ResBlock_1(x)
        x = self.mp(F.relu(self.Conv_2(x)))
        x = self.ResBlock_2(x)
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x


resNet = ResNet()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
resNet.to(device)
print(torch.cuda.is_available())
resNet.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(resNet.parameters(), lr=0.01, momentum=0.5)


# Set train cycle
def train(epoch):
    running_loss = 0.0
    for batch_idx, data in enumerate(train_dataloader, 0):
        inputs, target = data
        # Set GPU
        inputs, target = inputs.to(device), target.to(device)
        optimizer.zero_grad()
        # Forward + Backward + Update
        outputs = resNet(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if batch_idx % 300 == 299:
            print('[%d,%5d] loss: %.6f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0


def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_dataloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = resNet(images)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy on test set : %.3f' % (100 * correct / total))


if __name__ == "__main__":
    for epoch in range(30):
        train(epoch)
        test()
